#!/usr/bin/env python3

import socket
import os
import signal
import sys

line = 1

def outRecvData(sd, name='服务端'):
    global line

    data, addr = sd.recvfrom(1024)

    data = data.decode()

    if data == 'exit':
        return None

    print("\033[s\033[%d;1H\033[2K%s接受到(%s): %s\033[u"
                % (line, name, addr[0], data), end='', flush=True)
    line += 1
    if line == 11:
        line = 1

    return addr

def myExit(sig, stack):
    print("\033[15;1H\033[2K", end='', flush=True)
    os._exit(0)

def chatHandle(sd, addr, name="服务端"):
    if os.fork() == 0:
        while outRecvData(sd, name):
            pass
    else:
        while True:
            data = input("\033[13;1H\033[2K%s发送: " % name)
            sd.sendto(data.encode(), addr)
            if data == 'exit':
                break

    sd.close()
    os.kill(0, signal.SIGTERM)

def main():
    #  没有参数: 服务端 9000 0.0.0.0
    #  一个参数: 服务端 绑定端口号
    #  二个参数: 客户端 服务器口号 服务器IP
    argc = len(sys.argv)

    port = int(sys.argv[1]) if argc == 2 else 9000
    isSrv = True if argc < 3 else False
    name = "服务端" if isSrv else "客户端"
    ip = sys.argv[2] if argc == 3 else "0.0.0.0"

    print('\033[2J', end='', flush=True)
    signal.signal(signal.SIGTERM, myExit)
    sd = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sd.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    addr = (ip, port)

    if isSrv:
        sd.bind(addr)
        print("\033[1;1H\033[2K正在等待客户端发送数据.....", end='', flush=True)
        data, addr = sd.recvfrom(1024)
        sd.close()
        sd = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sd.sendto('请求通过'.encode('utf-8'), addr)
    else:
        sd.sendto('连接请求'.encode('utf-8'), addr)
        data, addr = sd.recvfrom(1024)

    chatHandle(sd, addr, name)

if __name__ == "__main__":
    main()
